import numpy as np
sbox = [[0x63, 0x7c, 0x77, 0x7b, 0xf2, 0x6b, 0x6f, 0xc5, 0x30, 0x01, 0x67, 0x2b, 0xfe, 0xd7, 0xab, 0x76],
        [0xca, 0x82, 0xc9, 0x7d, 0xfa, 0x59, 0x47, 0xf0, 0xad, 0xd4, 0xa2, 0xaf, 0x9c, 0xa4, 0x72, 0xc0],
        [0xb7, 0xfd, 0x93, 0x26, 0x36, 0x3f, 0xf7, 0xcc, 0x34, 0xa5, 0xe5, 0xf1, 0x71, 0xd8, 0x31, 0x15],
        [0x04, 0xc7, 0x23, 0xc3, 0x18, 0x96, 0x05, 0x9a, 0x07, 0x12, 0x80, 0xe2, 0xeb, 0x27, 0xb2, 0x75],
        [0x09, 0x83, 0x2c, 0x1a, 0x1b, 0x6e, 0x5a, 0xa0, 0x52, 0x3b, 0xd6, 0xb3, 0x29, 0xe3, 0x2f, 0x84],
        [0x53, 0xd1, 0x00, 0xed, 0x20, 0xfc, 0xb1, 0x5b, 0x6a, 0xcb, 0xbe, 0x39, 0x4a, 0x4c, 0x58, 0xcf],
        [0xd0, 0xef, 0xaa, 0xfb, 0x43, 0x4d, 0x33, 0x85, 0x45, 0xf9, 0x02, 0x7f, 0x50, 0x3c, 0x9f, 0xa8],
        [0x51, 0xa3, 0x40, 0x8f, 0x92, 0x9d, 0x38, 0xf5, 0xbc, 0xb6, 0xda, 0x21, 0x10, 0xff, 0xf3, 0xd2],
        [0xcd, 0x0c, 0x13, 0xec, 0x5f, 0x97, 0x44, 0x17, 0xc4, 0xa7, 0x7e, 0x3d, 0x64, 0x5d, 0x19, 0x73],
        [0x60, 0x81, 0x4f, 0xdc, 0x22, 0x2a, 0x90, 0x88, 0x46, 0xee, 0xb8, 0x14, 0xde, 0x5e, 0x0b, 0xdb],
        [0xe0, 0x32, 0x3a, 0x0a, 0x49, 0x06, 0x24, 0x5c, 0xc2, 0xd3, 0xac, 0x62, 0x91, 0x95, 0xe4, 0x79],
        [0xe7, 0xc8, 0x37, 0x6d, 0x8d, 0xd5, 0x4e, 0xa9, 0x6c, 0x56, 0xf4, 0xea, 0x65, 0x7a, 0xae, 0x08],
        [0xba, 0x78, 0x25, 0x2e, 0x1c, 0xa6, 0xb4, 0xc6, 0xe8, 0xdd, 0x74, 0x1f, 0x4b, 0xbd, 0x8b, 0x8a],
        [0x70, 0x3e, 0xb5, 0x66, 0x48, 0x03, 0xf6, 0x0e, 0x61, 0x35, 0x57, 0xb9, 0x86, 0xc1, 0x1d, 0x9e],
        [0xe1, 0xf8, 0x98, 0x11, 0x69, 0xd9, 0x8e, 0x94, 0x9b, 0x1e, 0x87, 0xe9, 0xce, 0x55, 0x28, 0xdf],
        [0x8c, 0xa1, 0x89, 0x0d, 0xbf, 0xe6, 0x42, 0x68, 0x41, 0x99, 0x2d, 0x0f, 0xb0, 0x54, 0xbb, 0x16] ]
Rconj = [[0x01, 0x00, 0x00, 0x00], [0x02, 0x00, 0x00, 0x00], [0x04, 0x00, 0x00, 0x00],[0x08, 0x00, 0x00, 0x00], [0x10, 0x00, 0x00, 0x00],
         [0x20, 0x00, 0x00, 0x00], [0x40, 0x00, 0x00, 0x00], [0x80, 0x00, 0x00, 0x00],[0x1B, 0x00, 0x00, 0x00], [0x36, 0x00, 0x00, 0x00]]

#AES密钥编排算法
def encode(Ini_w,size):
    w = np.mat(Ini_w).T
    for i in range(5,size+1):   #从第五列到第size列
        if (i) % 4 == 0:
            col_1 = w[:,i-4-1]
            col_2 = w[:,i-1-1]
            Rcon_i = int(i/4)-1                  #轮常量轮数
            col_2_T = T_fun(col_2,Rcon_i)   #T函数处理
            extend_col = xor_vector(col_1,col_2_T)
            w = np.insert(w,i-1,extend_col.T,axis=1)
        else:
            col_1 = w[:,i-4-1]
            col_2 = w[:,i-1-1]
            extend_col = xor_vector(col_1,col_2)  #异或操作
            w = np.insert(w,i-1,extend_col.T,axis=1) #列插入操作
    return w

#列向量异或函数
def xor_vector(col_1,col_2):     #参数为列向量，返回列向量
    result = col_1[0]^col_2[0]   #每一个对应元素分别异或
    for i in range(1,4):
        extend_temp = col_1[i]^col_2[i]
        result = np.insert(result,i,extend_temp,axis=0)
    return result
    
#T函数处理
def T_fun(col,Rcon_i):
    #循环左移（上移），，S盒替换，，轮常量异或
    temp = col[0]
    for i in range(3):
        col[i] = col[i+1]
    col[3] = temp
    #S盒替换
    for i in range(4):
        col[i]=S_substitution(col[i])

    #轮常量异或
    global Rconj
    Rconj_j = Rcon_i - 1      #索引对齐
    #将Rconj矩阵化，转置为列向量,,注意只第一次执行转置，后面的循环需要免疫保护
    if np.size(Rconj,0) != 4:
        Rconj = np.mat(Rconj).T
    result_col = xor_vector(col,Rconj[:,Rconj_j])


    return result_col

#S盒替换函数
def S_substitution(byte):           #byte是单元素矩阵，转为数字
    h_index = int(int(byte)/16)   #利用四位16，得到前后四位
    l_index = int(int(byte)%16)
    result = sbox[h_index][l_index]
    return result
#mat = np.insert(mat, 2, new_col.T, axis=1) # 新列插入第三列，记得转置插入列为行
#a = np.array([[1, 2], [3, 4], [5, 6]])




if __name__ == '__main__':
    w = []                #输入初始密钥矩阵
    for i in range(4):
        line = list(map(int,input('please input integers by lines: ').strip().split(',')))
        w.append(line)
    w = np.array(w)  #多维数组化

    w_New = encode(w,44) #AES，增加40个新列，共44列
    
    print(w_New)
    print('数出来为',np.size(w_New,1),'列')

'''    
1,2,3,4
5,6,7,8
1,6,7,9
2,0,4,8
for i in range(1,4):
    print(i)



'''





#list只在定义的时候分配空间，w = [……]按需求分配，对于需要增加的空间使用append方法
#   可以调用修改w[2] = 1，但是不能调用未创建的空间进行赋值
#input('please input one line: ')是一个字符串，是一个可索引字符的【原子单位】区别于列表。
#                                              区别于C语言的输入习惯
# 字符串.split('')是列表，map是处理列表的方法
#array多维数组化，（去掉了逗号，跟矩阵长得一样）
#num可以用16进制输入，但是存进去的是十进制
#w同理，可以用16进制输入，但是要加上0x这样就是int型
#python的除法默认float型，和C语言习惯不同，需要int()处理
#全局变量在引用的时候，要global 声明(局部初始化)
#矩阵列插入方法，插入列需要转置为行，不知道为什么